from enum import Enum

from ..config.config import GlobalConfiguration
from ..data_structures.base_structures import Map, HexCell, QuadCell
from .abs_quantifier import BaseQuantifier
from ..data_structures.base_structures import DemAttribute
from tqdm import tqdm
from statistics import stdev, mean
import h3

class TifQuantifierComponents(Enum):
    """量化tif的组件们"""

    "六角格网量化tif的组件们"
    QuantifyHexcellsWithElevation = 'quantify_hexcells_with_elevation'
    QuantifyHexcellsWithCurvature = 'quantify_hexcells_with_curvature'
    QuantifyHexcellsWithCv = 'quantify_hexcells_with_cv'
    QuantifyHexcellsWithExposure = 'quantify_hexcells_with_exposure'
    QuantifyHexcellsWithRelief = 'quantify_hexcells_with_relief'
    QuantifyHexcellsWithRoughness = 'quantify_hexcells_with_roughness'
    QuantifyHexcellsWithSlope = 'quantify_hexcells_with_slope'

    "四角格网量化tif的组件们"
    QuantifyQuadcellsWithElevation = 'quantify_quadcells_with_elevation'
    QuantifyQuadcellsWithCurvature = 'quantify_quadcells_with_curvature'
    QuantifyQuadcellsWithCv = 'quantify_quadcells_with_cv'
    QuantifyQuadcellsWithExposure = 'quantify_quadcells_with_exposure'
    QuantifyQuadcellsWithRelief = 'quantify_quadcells_with_relief'
    QuantifyQuadcellsWithRoughness = 'quantify_quadcells_with_roughness'
    QuantifyQuadcellsWithSlope = 'quantify_quadcells_with_slope'

    @staticmethod
    def quantify_hexcells_with_elevation(data, transform, nodata_value, map:Map):
        """量化高程"""
        for hexcell in tqdm(map.cells.values(), desc="量化高程:"):
            hexcell:HexCell
            # 初始化顶点高程数组
            vertex_elevation = []  

            # 填充顶点高程数组
            for lat, lon in hexcell.vertices:               
                elevation = BaseQuantifier.get_value(data, transform, lon, lat, nodata_value)
                vertex_elevation.append(elevation)

            # 计算中心高程
            count = 0
            sum_elevation = 0
            for elevation in vertex_elevation:
                if elevation is not None:
                    count += 1
                    sum_elevation += elevation
            center_elevation = sum_elevation / count if count > 0 else None

            # 填充cell的属性数组
            hexcell.attributes.append(DemAttribute.Elevation(center_elevation))

    @staticmethod
    def quantify_hexcells_with_curvature(data, transform, nodata_value, map:Map):
        """领域法量化曲率"""
        # 估算单元格大小（基于H3分辨率）
        sample_cell = next(iter(map.cells.values()))
        import h3
        resolution = h3.h3_get_resolution(sample_cell.index)
        cell_size = h3.edge_length(resolution, unit='m')
        
        for cell in tqdm(map.cells.values(), desc=f"量化曲率:"):
            cell: HexCell
            cell_elevation = cell.get_attribute(DemAttribute.Elevation).get_value()
            if cell_elevation is None:
                cell.attributes.append(DemAttribute.Curvature(None))
                continue

            # 获取3x3邻域的高程值
            elevations = BaseQuantifier.get_neighbor_elevations(map, cell)

            curvature_value = BaseQuantifier.calculate_curvature(elevations, cell_size)
            cell.attributes.append(DemAttribute.Curvature(curvature_value))
    
    @staticmethod
    def quantify_hexcells_with_cv(data, transform, nodata_value, map:Map):
        """量化高程变异系数"""
        for cell in tqdm(map.cells.values(), desc="量化高程变异系数:"):
            cell:HexCell
            elevations = []
            cell_elevation = cell.get_attribute(DemAttribute.Elevation).get_value()
            # 包含自身
            if cell_elevation is not None:
                elevations.append(cell_elevation)

            # 包含6个邻接单元
            for neighbor_index in cell.neighbors:
                neighbor_cell = map.cells.get(neighbor_index)
                neighbor_cell:HexCell
                if not neighbor_cell:
                    continue
                neighbor_cell_elevation = neighbor_cell.get_attribute(DemAttribute.Elevation).get_value()
                if neighbor_cell and neighbor_cell_elevation is not None:
                    elevations.append(neighbor_cell_elevation)

            if len(elevations) >= 2:
                avg = mean(elevations)    # 平均值
                sd = stdev(elevations)     # 标准差
                cv = DemAttribute.ElevationCoefficientOfVariation(round(sd / avg, 4) if avg != 0 else None)
                cell.attributes.append(cv)
            else:
                cell.attributes.append(DemAttribute.ElevationCoefficientOfVariation(None))
        return map
    
    @staticmethod
    def quantify_hexcells_with_exposure(data, transform, nodata_value, map:Map):
        """邻域法量化坡向"""
        # 估算单元格大小（基于H3分辨率）
        sample_cell = next(iter(map.cells.values()))
        import h3
        resolution = h3.h3_get_resolution(sample_cell.index)
        cell_size = h3.edge_length(resolution, unit='m')
        
        for cell in tqdm(map.cells.values(), desc="量化坡向: "):
            cell_elevation = cell.get_attribute(DemAttribute.Elevation).get_value()
            if cell_elevation is None:
                cell.attributes.append(DemAttribute.Exposure(None))
                continue
            
            # 获取3x3邻域的高程值
            elevations = BaseQuantifier.get_neighbor_elevations(map, cell)
            
            # 计算梯度
            dz_dx, dz_dy = BaseQuantifier.calculate_gradient(elevations, cell_size)
            
            if dz_dx is not None and dz_dy is not None:
                # 计算坡向
                exposure_value = BaseQuantifier.calculate_exposure_from_gradient(dz_dx, dz_dy)
            else:
                exposure_value = None
                
            cell.attributes.append(DemAttribute.Exposure(exposure_value))
    
    @staticmethod
    def quantify_hexcells_with_relief(data, transform, nodata_value, map:Map):
        """邻域法量化地形起伏度"""
        for cell in tqdm(map.cells.values(), desc="量化地形起伏度: "):
            cell:HexCell
            elevations = []
            cell_elevation = cell.get_attribute(DemAttribute.Elevation).get_value()
            # 包含自身
            if cell_elevation is not None:
                elevations.append(cell_elevation)
            # 包含6个邻接单元
            for neighbor_index in cell.neighbors:
                neighbor_cell = map.cells.get(neighbor_index)
                if not neighbor_cell:
                    continue
                neighbor_cell_elevation = neighbor_cell.get_attribute(DemAttribute.Elevation).get_value()
                if neighbor_cell and neighbor_cell_elevation is not None:
                    elevations.append(neighbor_cell_elevation)
            # 计算地形起伏度
            if len(elevations) >= 2:
                relief = DemAttribute.Relief(max(elevations) - min(elevations))
                cell.attributes.append(relief)
            else:
                relief = DemAttribute.Relief(0)
                cell.attributes.append(relief)
        return map
    
    @staticmethod
    def quantify_hexcells_with_roughness(data, transform, nodata_value, map:Map):
        """邻域法量化地形粗糙度"""
        for cell in tqdm(map.cells.values(), desc="量化地形粗糙度:"):
            cell:HexCell
            elevations = []
            cell_elevation = cell.get_attribute(DemAttribute.Elevation).get_value()
            # 包含自身
            if cell_elevation is not None:
                elevations.append(cell_elevation)
            # 包含6个邻接单元
            for neighbor_index in cell.neighbors:
                neighbor_cell = map.cells.get(neighbor_index)
                if not neighbor_cell:
                    continue
                neighbor_cell_elevation = neighbor_cell.get_attribute(DemAttribute.Elevation).get_value()
                if neighbor_cell and neighbor_cell_elevation is not None:
                    elevations.append(neighbor_cell_elevation)
            # 计算地形粗糙度
            if len(elevations) >= 2:
                sd = stdev(elevations)
                roughness = DemAttribute.Roughness(sd)
                cell.attributes.append(roughness)
            else:
                roughness = DemAttribute.Roughness(0)
                cell.attributes.append(roughness)
        return map
    
    @staticmethod
    def quantify_hexcells_with_slope(data, transform, nodata_value, map:Map):
        for hexcell in tqdm(map.cells.values(), desc="量化坡度:"):
            hexcell:HexCell
            # 初始化顶点高程数组
            vertex_elevation = []  

            # 填充顶点高程数组
            for lat, lon in hexcell.vertices:               
                elevation = BaseQuantifier.get_value(data, transform, lon, lat, nodata_value)
                vertex_elevation.append(elevation)

            # 计算坡度
            hexcell_index = hexcell.index
            hexcell_size = h3.edge_length(h3.h3_get_resolution(hexcell_index), unit='m')
            center_elevation = hexcell.get_attribute(DemAttribute.Elevation).get_value()
            if center_elevation == None:
                continue
            center_slope = BaseQuantifier.calculate_center_slope(hexcell_size, vertex_elevation, center_elevation, grid_type='hex')

            # 填充cell的属性数组
            hexcell.attributes.append(DemAttribute.Slope(center_slope))

    @staticmethod
    def quantify_quadcells_with_elevation(data, transform, nodata_value, map:Map):
        """量化四角格网高程"""
        for quadcell in tqdm(map.cells.values(), desc="量化高程:"):
            quadcell:QuadCell
            # 初始化顶点高程数组
            vertex_elevation = []

            # 填充顶点高程数组
            for lat, lon in quadcell.vertices:
                elevation = BaseQuantifier.get_value(data, transform, lon, lat, nodata_value)
                vertex_elevation.append(elevation)

            # 计算中心高程
            count = 0
            sum_elevation = 0
            for elevation in vertex_elevation:
                if elevation is not None:
                    count += 1
                    sum_elevation += elevation
            center_elevation = sum_elevation / count if count > 0 else None

            # 填充cell的属性数组
            quadcell.attributes.append(DemAttribute.Elevation(center_elevation))

    @staticmethod
    def quantify_quadcells_with_curvature(data, transform, nodata_value, map: Map):
        """领域法量化四角格网曲率"""
        precision = GlobalConfiguration.geohash_resolution
        cell_size = BaseQuantifier.estimate_cell_size(precision)  # 获取geohash的格网边长

        for cell in tqdm(map.cells.values(), desc=f"量化曲率（Geohash）:"):
            cell:QuadCell
            cell_elevation = cell.get_attribute(DemAttribute.Elevation).get_value()
            if cell_elevation is None:
                cell.attributes.append(DemAttribute.Curvature(None))
                continue

            # 获取3x3邻域的高程值
            elevations = BaseQuantifier.get_neighbor_elevations(map, cell)
            curvature_value = BaseQuantifier.calculate_curvature(elevations, cell_size)
            cell.attributes.append(DemAttribute.Curvature(curvature_value))

    @staticmethod
    def quantify_quadcells_with_cv(data, transform, nodata_value, map: Map):
        """量化四角格网高程变异系数"""
        for cell in tqdm(map.cells.values(), desc="量化高程变异系数:"):
            cell: QuadCell
            elevations = []
            cell_elevation = cell.get_attribute(DemAttribute.Elevation).get_value()
            # 包含自身
            if cell_elevation is not None:
                elevations.append(cell_elevation)

            # 包含6个邻接单元
            for neighbor_index in cell.neighbors:
                neighbor_cell = map.cells.get(neighbor_index)
                neighbor_cell: QuadCell
                if not neighbor_cell:
                    continue
                neighbor_cell_elevation = neighbor_cell.get_attribute(DemAttribute.Elevation).get_value()
                if neighbor_cell and neighbor_cell_elevation is not None:
                    elevations.append(neighbor_cell_elevation)

            if len(elevations) >= 2:
                avg = mean(elevations)  # 平均值
                sd = stdev(elevations)  # 标准差
                cv = DemAttribute.ElevationCoefficientOfVariation(round(sd / avg, 4) if avg != 0 else None)
                cell.attributes.append(cv)
            else:
                cell.attributes.append(DemAttribute.ElevationCoefficientOfVariation(None))
        return map

    @staticmethod
    def quantify_quadcells_with_exposure(data, transform, nodata_value, map: Map):
        """邻域法量化四角格网坡向"""

        """邻域法量化坡向"""
        # 估算单元格大小（基于H3分辨率）
        precision = GlobalConfiguration.geohash_resolution
        cell_size = BaseQuantifier.estimate_cell_size(precision)  # 获取geohash的格网边长

        for cell in tqdm(map.cells.values(), desc="量化坡向: "):
            cell_elevation = cell.get_attribute(DemAttribute.Elevation).get_value()
            if cell_elevation is None:
                cell.attributes.append(DemAttribute.Exposure(None))
                continue

            # 获取3x3邻域的高程值
            elevations = BaseQuantifier.get_neighbor_elevations(map, cell)

            # 计算梯度
            dz_dx, dz_dy = BaseQuantifier.calculate_gradient(elevations, cell_size)

            if dz_dx is not None and dz_dy is not None:
                # 计算坡向
                exposure_value = BaseQuantifier.calculate_exposure_from_gradient(dz_dx, dz_dy)
            else:
                exposure_value = None

            cell.attributes.append(DemAttribute.Exposure(exposure_value))

    @staticmethod
    def quantify_quadcells_with_relief(data, transform, nodata_value, map: Map):
        """邻域法量化四角格网地形起伏度"""
        for cell in tqdm(map.cells.values(), desc="量化地形起伏度: "):
            cell:QuadCell
            elevations = []
            cell_elevation = cell.get_attribute(DemAttribute.Elevation).get_value()
            # 包含自身
            if cell_elevation is not None:
                elevations.append(cell_elevation)
            # 包含8个邻接单元
            for neighbor_index in cell.neighbors:
                neighbor_cell = map.cells.get(neighbor_index)
                if not neighbor_cell:
                    continue
                neighbor_cell_elevation = neighbor_cell.get_attribute(DemAttribute.Elevation).get_value()
                if neighbor_cell and neighbor_cell_elevation is not None:
                    elevations.append(neighbor_cell_elevation)
            # 计算地形起伏度
            if len(elevations) >= 2:
                relief = DemAttribute.Relief(max(elevations) - min(elevations))
                cell.attributes.append(relief)
            else:
                relief = DemAttribute.Relief(0)
                cell.attributes.append(relief)
        return map

    @staticmethod
    def quantify_quadcells_with_roughness(data, transform, nodata_value, map: Map):
        """邻域法量化地形粗糙度"""
        for cell in tqdm(map.cells.values(), desc="量化地形粗糙度:"):
            cell: QuadCell
            elevations = []
            cell_elevation = cell.get_attribute(DemAttribute.Elevation).get_value()
            # 包含自身
            if cell_elevation is not None:
                elevations.append(cell_elevation)
            # 包含6个邻接单元
            for neighbor_index in cell.neighbors:
                neighbor_cell = map.cells.get(neighbor_index)
                if not neighbor_cell:
                    continue
                neighbor_cell_elevation = neighbor_cell.get_attribute(DemAttribute.Elevation).get_value()
                if neighbor_cell and neighbor_cell_elevation is not None:
                    elevations.append(neighbor_cell_elevation)
            # 计算地形粗糙度
            if len(elevations) >= 2:
                sd = stdev(elevations)
                roughness = DemAttribute.Roughness(sd)
                cell.attributes.append(roughness)
            else:
                roughness = DemAttribute.Roughness(0)
                cell.attributes.append(roughness)
        return map

    @staticmethod
    def quantify_quadcells_with_slope(data, transform, nodata_value, map: Map):
        for quadcell in tqdm(map.cells.values(), desc="量化坡度:"):
            quadcell:QuadCell
            # 初始化顶点高程数组
            vertex_elevation = []

            # 填充顶点高程数组
            for lat, lon in quadcell.vertices:
                elevation = BaseQuantifier.get_value(data, transform, lon, lat, nodata_value)
                vertex_elevation.append(elevation)

            # 计算坡度
            precision = GlobalConfiguration.geohash_resolution
            quadcell_size = BaseQuantifier.estimate_cell_size(precision)
            center_elevation = quadcell.get_attribute(DemAttribute.Elevation).get_value()
            if center_elevation == None:
                continue
            center_slope = BaseQuantifier.calculate_center_slope(quadcell_size, vertex_elevation, center_elevation, grid_type='quad')

            # 填充cell的属性数组
            quadcell.attributes.append(DemAttribute.Slope(center_slope))